// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <memory>
#include <string>
#include <vector>

#include <ktest/ktest.h>

#include <nebula/array/util.h>
#include <nebula/compute/kernel.h>
#include <turbo/utility/status.h>
#include <nebula/testing/ktest_util.h>
#include <nebula/types/type.h>
#include <nebula/util/key_value_metadata.h>

namespace nebula::compute {

// ----------------------------------------------------------------------
// TypeMatcher

    TEST(TypeMatcher, SameTypeId) {
        std::shared_ptr<TypeMatcher> matcher = match::SameTypeId(Type::DECIMAL);
        ASSERT_TRUE(matcher->matches(*decimal(12, 2)));
        ASSERT_FALSE(matcher->matches(*int8()));

        ASSERT_EQ("Type::DECIMAL128", matcher->to_string());

        ASSERT_TRUE(matcher->equals(*matcher));
        ASSERT_TRUE(matcher->equals(*match::SameTypeId(Type::DECIMAL)));
        ASSERT_FALSE(matcher->equals(*match::SameTypeId(Type::TIMESTAMP)));
    }

    TEST(TypeMatcher, TimestampTypeUnit) {
        auto matcher = match::TimestampTypeUnit(TimeUnit::MILLI);
        auto matcher2 = match::Time32TypeUnit(TimeUnit::MILLI);

        ASSERT_TRUE(matcher->matches(*timestamp(TimeUnit::MILLI)));
        ASSERT_TRUE(matcher->matches(*timestamp(TimeUnit::MILLI, "utc")));
        ASSERT_FALSE(matcher->matches(*timestamp(TimeUnit::SECOND)));
        ASSERT_FALSE(matcher->matches(*time32(TimeUnit::MILLI)));
        ASSERT_TRUE(matcher2->matches(*time32(TimeUnit::MILLI)));

        // Check ToString representation
        ASSERT_EQ("timestamp(s)", match::TimestampTypeUnit(TimeUnit::SECOND)->to_string());
        ASSERT_EQ("timestamp(ms)", match::TimestampTypeUnit(TimeUnit::MILLI)->to_string());
        ASSERT_EQ("timestamp(us)", match::TimestampTypeUnit(TimeUnit::MICRO)->to_string());
        ASSERT_EQ("timestamp(ns)", match::TimestampTypeUnit(TimeUnit::NANO)->to_string());

        // Equals implementation
        ASSERT_TRUE(matcher->equals(*matcher));
        ASSERT_TRUE(matcher->equals(*match::TimestampTypeUnit(TimeUnit::MILLI)));
        ASSERT_FALSE(matcher->equals(*match::TimestampTypeUnit(TimeUnit::MICRO)));
        ASSERT_FALSE(matcher->equals(*match::Time32TypeUnit(TimeUnit::MILLI)));
    }

    TEST(TypeMatcher, RunEndInteger) {
        auto matcher = match::RunEndInteger();
        ASSERT_FALSE(matcher->matches(*int8()));
        ASSERT_FALSE(matcher->matches(*uint8()));
        ASSERT_FALSE(matcher->matches(*uint16()));
        ASSERT_FALSE(matcher->matches(*uint32()));
        ASSERT_FALSE(matcher->matches(*uint64()));

        ASSERT_TRUE(matcher->matches(*int16()));
        ASSERT_TRUE(matcher->matches(*int32()));
        ASSERT_TRUE(matcher->matches(*int64()));
    }

    TEST(TypeMatcher, RunEndEncoded) {
        auto string_ree_matcher = match::RunEndEncoded(match::SameTypeId(Type::STRING));
        auto int32_ree_matcher = match::RunEndEncoded(match::SameTypeId(Type::INT32));
        auto empty_ree_matcher =
                match::RunEndEncoded(match::SameTypeId(Type::INT8), match::SameTypeId(Type::INT32));
        for (auto run_end_type: {int16(), int32(), int64()}) {
            ASSERT_TRUE(string_ree_matcher->matches(*run_end_encoded(run_end_type, utf8())));
            ASSERT_FALSE(string_ree_matcher->matches(*run_end_encoded(run_end_type, int32())));

            ASSERT_FALSE(int32_ree_matcher->matches(*run_end_encoded(run_end_type, utf8())));
            ASSERT_TRUE(int32_ree_matcher->matches(*run_end_encoded(run_end_type, int32())));

            ASSERT_FALSE(empty_ree_matcher->matches(*run_end_encoded(run_end_type, int32())));
        }
    }

// ----------------------------------------------------------------------
// InputType

    TEST(InputType, AnyTypeConstructor) {
        // Check the ANY_TYPE ctors
        InputType ty;
        ASSERT_EQ(InputType::ANY_TYPE, ty.kind());
    }

    TEST(InputType, Constructors) {
        // Exact type constructor
        InputType ty1(int8());
        ASSERT_EQ(InputType::EXACT_TYPE, ty1.kind());
        AssertTypeEqual(*int8(), *ty1.type());

        InputType ty1_implicit = int8();
        ASSERT_TRUE(ty1.equals(ty1_implicit));

        // Same type id constructor
        InputType ty2(Type::DECIMAL);
        ASSERT_EQ(InputType::USE_TYPE_MATCHER, ty2.kind());
        ASSERT_EQ("Type::DECIMAL128", ty2.to_string());
        ASSERT_TRUE(ty2.type_matcher().matches(*decimal(12, 2)));
        ASSERT_FALSE(ty2.type_matcher().matches(*int16()));

        // implicit construction in a vector
        std::vector<InputType> types = {int8(), InputType(Type::DECIMAL)};
        ASSERT_TRUE(types[0].equals(ty1));
        ASSERT_TRUE(types[1].equals(ty2));

        // Copy constructor
        InputType ty3 = ty1;
        InputType ty4 = ty2;
        ASSERT_TRUE(ty3.equals(ty1));
        ASSERT_TRUE(ty4.equals(ty2));

        // Move constructor
        InputType ty5 = std::move(ty3);
        InputType ty6 = std::move(ty4);
        ASSERT_TRUE(ty5.equals(ty1));
        ASSERT_TRUE(ty6.equals(ty2));

        // ToString
        ASSERT_EQ("int8", ty1.to_string());
        ASSERT_EQ("Type::DECIMAL128", ty2.to_string());

        InputType ty7(match::TimestampTypeUnit(TimeUnit::MICRO));
        ASSERT_EQ("timestamp(us)", ty7.to_string());

        InputType ty8;
        ASSERT_EQ("any", ty8.to_string());
    }

    TEST(InputType, Equals) {
        InputType t1 = int8();
        InputType t2 = int8();
        InputType t3 = int32();

        InputType t5(Type::DECIMAL);
        InputType t6(Type::DECIMAL);

        ASSERT_TRUE(t1.equals(t2));
        ASSERT_EQ(t1, t2);
        ASSERT_NE(t1, t3);

        ASSERT_FALSE(t1.equals(t5));
        ASSERT_NE(t1, t5);

        ASSERT_EQ(t5, t5);
        ASSERT_EQ(t5, t6);

        // NOTE: For the time being, we treat int32() and Type::INT32 as being
        // different. This could obviously be fixed later to make these equivalent
        ASSERT_NE(InputType(int8()), InputType(Type::INT32));

        // Check that field metadata excluded from equality checks
        InputType t9 = list(
                field("item", utf8(), /*nullable=*/true, key_value_metadata({"foo"}, {"bar"})));
        InputType t10 = list(field("item", utf8()));
        ASSERT_TRUE(t9.equals(t10));
    }

    TEST(InputType, Hash) {
        InputType t0;
        InputType t1 = int8();
        InputType t2(Type::DECIMAL);

        // These checks try to determine first of all whether Hash always returns the
        // same value, and whether the elements of the type are all incorporated into
        // the Hash
        ASSERT_EQ(t0.hash(), t0.hash());
        ASSERT_EQ(t1.hash(), t1.hash());
        ASSERT_EQ(t2.hash(), t2.hash());
        ASSERT_NE(t0.hash(), t1.hash());
        ASSERT_NE(t0.hash(), t2.hash());
        ASSERT_NE(t1.hash(), t2.hash());
    }

    TEST(InputType, matches) {
        InputType input1 = int8();

        ASSERT_TRUE(input1.matches(*int8()));
        ASSERT_TRUE(input1.matches(*int8()));
        ASSERT_FALSE(input1.matches(*int16()));

        InputType input2(Type::DECIMAL);
        ASSERT_TRUE(input2.matches(*decimal(12, 2)));

        auto ty2 = decimal(12, 2);
        auto ty3 = float64();
        ASSERT_OK_AND_ASSIGN(std::shared_ptr<Array> arr2, MakeArrayOfNull(ty2, 1));
        ASSERT_OK_AND_ASSIGN(std::shared_ptr<Array> arr3, MakeArrayOfNull(ty3, 1));
        ASSERT_OK_AND_ASSIGN(std::shared_ptr<Scalar> scalar2, arr2->get_scalar(0));
        ASSERT_TRUE(input2.matches(Datum(arr2)));
        ASSERT_TRUE(input2.matches(Datum(scalar2)));
        ASSERT_FALSE(input2.matches(*ty3));
        ASSERT_FALSE(input2.matches(arr3));
    }

// ----------------------------------------------------------------------
// OutputType

    TEST(OutputType, Constructors) {
        OutputType ty1 = int8();
        ASSERT_EQ(OutputType::FIXED, ty1.kind());
        AssertTypeEqual(*int8(), *ty1.type());

        auto DummyResolver = [](KernelContext *,
                                const std::vector<TypeHolder> &args) -> turbo::Result<TypeHolder> {
            return int32();
        };
        OutputType ty2(DummyResolver);
        ASSERT_EQ(OutputType::COMPUTED, ty2.kind());

        ASSERT_OK_AND_ASSIGN(TypeHolder out_type2, ty2.resolve(nullptr, {}));
        ASSERT_EQ(out_type2, int32());

        // Copy constructor
        OutputType ty3 = ty1;
        ASSERT_EQ(OutputType::FIXED, ty3.kind());
        AssertTypeEqual(*ty1.type(), *ty3.type());

        OutputType ty4 = ty2;
        ASSERT_EQ(OutputType::COMPUTED, ty4.kind());
        ASSERT_OK_AND_ASSIGN(TypeHolder out_type4, ty4.resolve(nullptr, {}));
        ASSERT_EQ(out_type4, int32());

        // Move constructor
        OutputType ty5 = std::move(ty1);
        ASSERT_EQ(OutputType::FIXED, ty5.kind());
        AssertTypeEqual(*int8(), *ty5.type());

        OutputType ty6 = std::move(ty4);
        ASSERT_EQ(OutputType::COMPUTED, ty6.kind());
        ASSERT_OK_AND_ASSIGN(TypeHolder out_type6, ty6.resolve(nullptr, {}));
        ASSERT_EQ(out_type6, int32());

        // ToString

        // ty1 was copied to ty3
        ASSERT_EQ("int8", ty3.to_string());
        ASSERT_EQ("computed", ty2.to_string());
    }

    TEST(OutputType, resolve) {
        OutputType ty1(int32());

        ASSERT_OK_AND_ASSIGN(TypeHolder result, ty1.resolve(nullptr, {}));
        ASSERT_EQ(result, int32());

        ASSERT_OK_AND_ASSIGN(result, ty1.resolve(nullptr, {int8()}));
        ASSERT_EQ(result, int32());

        ASSERT_OK_AND_ASSIGN(result, ty1.resolve(nullptr, {int8(), int8()}));
        ASSERT_EQ(result, int32());

        auto resolver = [](KernelContext *,
                           const std::vector<TypeHolder> &args) -> turbo::Result<TypeHolder> {
            return args[0];
        };
        OutputType ty2(resolver);

        ASSERT_OK_AND_ASSIGN(result, ty2.resolve(nullptr, {utf8()}));
        ASSERT_EQ(result, utf8());

        // Type resolver that returns an error
        OutputType ty3(
                [](KernelContext *ctx, const std::vector<TypeHolder> &types) -> turbo::Result<TypeHolder> {
                    // NB: checking the value types versus the function arity should be
                    // validated elsewhere, so this is just for illustration purposes
                    if (types.size() == 0) {
                        return turbo::invalid_argument_error("Need at least one argument");
                    }
                    return types[0];
                });
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, ty3.resolve(nullptr, {}));

        // Type resolver that returns a fixed value
        OutputType ty4(
                [](KernelContext *ctx, const std::vector<TypeHolder> &types) -> turbo::Result<TypeHolder> {
                    return int32();
                });

        ASSERT_OK_AND_ASSIGN(result, ty4.resolve(nullptr, {int8()}));
        ASSERT_EQ(result, int32());
        ASSERT_OK_AND_ASSIGN(result, ty4.resolve(nullptr, {int8()}));
        ASSERT_EQ(result, int32());
    }

// ----------------------------------------------------------------------
// KernelSignature

    TEST(KernelSignature, Basics) {
        // (int8, decimal) -> utf8
        std::vector<InputType> in_types({int8(), InputType(Type::DECIMAL)});
        OutputType out_type(utf8());

        KernelSignature sig(in_types, out_type);
        ASSERT_EQ(2, sig.in_types().size());
        ASSERT_TRUE(sig.in_types()[0].type()->equals(*int8()));
        ASSERT_TRUE(sig.in_types()[0].matches(*int8()));
        ASSERT_TRUE(sig.in_types()[1].matches(*decimal(12, 2)));
    }

    TEST(KernelSignature, Equals) {
        KernelSignature sig1({}, utf8());
        KernelSignature sig1_copy({}, utf8());
        KernelSignature sig2({int8()}, utf8());

        // Output type doesn't matter (for now)
        KernelSignature sig3({int8()}, int32());

        KernelSignature sig4({int8(), int16()}, utf8());
        KernelSignature sig4_copy({int8(), int16()}, utf8());
        KernelSignature sig5({int8(), int16(), int32()}, utf8());

        ASSERT_EQ(sig1, sig1);

        ASSERT_EQ(sig2, sig3);
        ASSERT_NE(sig3, sig4);

        // Different sig objects, but same sig
        ASSERT_EQ(sig1, sig1_copy);
        ASSERT_EQ(sig4, sig4_copy);

        // Match first 2 args, but not third
        ASSERT_NE(sig4, sig5);
    }

    TEST(KernelSignature, VarArgsEquals) {
        KernelSignature sig1({int8()}, utf8(), /*is_varargs=*/true);
        KernelSignature sig2({int8()}, utf8(), /*is_varargs=*/true);
        KernelSignature sig3({int8()}, utf8());

        ASSERT_EQ(sig1, sig2);
        ASSERT_NE(sig2, sig3);
    }

    TEST(KernelSignature, Hash) {
        // Some basic tests to ensure that the hashes are deterministic and that all
        // input arguments are incorporated
        KernelSignature sig1({}, utf8());
        KernelSignature sig2({int8()}, utf8());
        KernelSignature sig3({int8(), int32()}, utf8());

        ASSERT_EQ(sig1.hash(), sig1.hash());
        ASSERT_EQ(sig2.hash(), sig2.hash());
        ASSERT_NE(sig1.hash(), sig2.hash());
        ASSERT_NE(sig2.hash(), sig3.hash());
    }

    TEST(KernelSignature, matches_inputs) {
        // () -> boolean
        KernelSignature sig1({}, boolean());

        ASSERT_TRUE(sig1.matches_inputs({}));
        ASSERT_FALSE(sig1.matches_inputs({int8()}));

        // (int8, decimal) -> boolean
        KernelSignature sig2({int8(), InputType(Type::DECIMAL)}, boolean());

        ASSERT_FALSE(sig2.matches_inputs({}));
        ASSERT_FALSE(sig2.matches_inputs({int8()}));
        ASSERT_TRUE(sig2.matches_inputs({int8(), decimal(12, 2)}));

        // (int8, int32) -> boolean
        KernelSignature sig3({int8(), int32()}, boolean());

        ASSERT_FALSE(sig3.matches_inputs({}));

        // Unqualified, these are ANY type and do not match because the kernel
        // requires a scalar and an array
        ASSERT_TRUE(sig3.matches_inputs({int8(), int32()}));
        ASSERT_FALSE(sig3.matches_inputs({int8(), int16()}));
    }

    TEST(KernelSignature, VarArgsMatchesInputs) {
        {
            KernelSignature sig({int8()}, utf8(), /*is_varargs=*/true);

            std::vector<TypeHolder> args = {int8()};
            ASSERT_TRUE(sig.matches_inputs(args));
            args.push_back(int8());
            args.push_back(int8());
            ASSERT_TRUE(sig.matches_inputs(args));
            args.push_back(int32());
            ASSERT_FALSE(sig.matches_inputs(args));
        }
        {
            KernelSignature sig({int8(), utf8()}, utf8(), /*is_varargs=*/true);

            std::vector<TypeHolder> args = {int8()};
            ASSERT_TRUE(sig.matches_inputs(args));
            args.push_back(utf8());
            args.push_back(utf8());
            ASSERT_TRUE(sig.matches_inputs(args));
            args.push_back(int32());
            ASSERT_FALSE(sig.matches_inputs(args));
        }
    }

    TEST(KernelSignature, ToString) {
        std::vector<InputType> in_types = {InputType(int8()), InputType(Type::DECIMAL),
                                           InputType(utf8())};
        KernelSignature sig(in_types, utf8());
        ASSERT_EQ("(int8, Type::DECIMAL128, string) -> string", sig.to_string());

        OutputType out_type(
                [](KernelContext *, const std::vector<TypeHolder> &args) -> turbo::Result<TypeHolder> {
                    return turbo::invalid_argument_error("NYI");
                });
        KernelSignature sig2({int8(), Type::DECIMAL}, out_type);
        ASSERT_EQ("(int8, Type::DECIMAL128) -> computed", sig2.to_string());
    }

    TEST(KernelSignature, VarArgsToString) {
        KernelSignature sig({int8()}, utf8(), /*is_varargs=*/true);
        ASSERT_EQ("varargs[int8*] -> string", sig.to_string());

        KernelSignature sig2({utf8(), int8()}, utf8(), /*is_varargs=*/true);
        ASSERT_EQ("varargs[string, int8*] -> string", sig2.to_string());
    }

}  // namespace nebula::compute
